#include <iostream>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>


static void _warp_polar(const cv::Mat& src, cv::Mat& dst, cv::Point2f center, double maxRadius, bool semiLog=false) {
    int dwidth = cvRound(maxRadius);
    int dheight = cvRound(maxRadius * CV_PI);
    dst.create(dheight, dwidth, CV_8UC1);
    
    double kangle = dheight / (CV_PI * 2);
    double kmag = dwidth / maxRadius;
    
    if (semiLog) {
        kmag = dwidth / std::log(maxRadius);
    }
    
    for (int phi = 0; phi < dheight; phi++) {
        for (int rho = 0; rho < dwidth; rho++) {
            float frho = (float)(rho / kmag);
            
            if (semiLog) {
                frho = (float)(std::exp(rho / kmag) - 1.0);
            }
            
            double angle = phi / kangle;
            double x = frho * std::cos(angle);
            double y = frho * std::sin(angle);
            int src_x = cv::saturate_cast<int>(x + center.x);
            int src_y = cv::saturate_cast<int>(y + center.y);
            if (src_x >= 0 && src_x < src.cols && src_y >= 0 && src_y < src.rows) {
                dst.ptr<uchar>(phi)[rho] = src.ptr<uchar>(src_y)[src_x];
            } else {
                dst.ptr<uchar>(phi)[rho] = 0;
            }
        }
    }
}

static void _warp_polar_inverse(const cv::Mat& src, cv::Mat& dst, int dwidth, int dheight, cv::Point2f center, double maxRadius, bool semiLog=false) {
    dst.create(dheight, dwidth, CV_8UC1);
    cv::Mat _src;
    cv::copyMakeBorder(src, _src, 1, 1, 0, 0, cv::BORDER_WRAP);
    double kmag = src.cols / maxRadius;
    double kangle = src.rows / (CV_PI * 2);
    if (semiLog) {
        kmag = src.cols / std::log(maxRadius);
    }
    
    cv::Mat mat_dx(1, 1, CV_32FC1), mat_dy(1, 1, CV_32FC1);
    cv::Mat mat_magnitude(1, 1, CV_32FC1), mat_angle(1, 1, CV_32FC1);
    
    for (int y = 0; y < dheight; y++) {
        for (int x = 0; x < dwidth; x++) {
            float dx = (float)x - center.x;
            float dy = (float)y - center.y;
            
            mat_dx.ptr<float>(0)[0] = dx;
            mat_dy.ptr<float>(0)[0] = dy;
            cv::cartToPolar(mat_dx, mat_dy, mat_magnitude, mat_angle, false);
            if (semiLog) {
                mat_magnitude.ptr<float>(0)[0] = std::log(mat_magnitude.ptr<float>(0)[0]+1.0);
            }
            
            double rho = mat_magnitude.ptr<float>(0)[0] * kmag;
            double phi = mat_angle.ptr<float>(0)[0] * kangle;
            
            float fmapx = (float)rho;
            float fmapy = (float)phi + 1.0;
            
            int src_x = cv::saturate_cast<int>(fmapx);
            int src_y = cv::saturate_cast<int>(fmapy);
            if (src_x >= 0 && src_x < _src.cols && src_y >= 0 && src_y < _src.rows) {
                dst.ptr<uchar>(y)[x] = _src.ptr<uchar>(src_y)[src_x];
            } else {
                dst.ptr<uchar>(y)[x] = 0;
            }
        }
    }
}

int main(int argc, char **argv) {
    cv::Mat src = cv::imread("/home/xlll/Downloads/opencv/samples/data/lena.jpg", cv::IMREAD_GRAYSCALE);
    if (src.empty()) {
        std::cout << "failed to read image" << std::endl;
        return EXIT_FAILURE;
    }
    
    cv::Mat linearPolar1, linearPolar2;
    int flags = cv::INTER_NEAREST;
    cv::Point2f center((float)src.cols / 2, (float)src.rows / 2);
    double maxRadius = 0.7 * std::min(center.y, center.x);
    
    cv::warpPolar(src, linearPolar1, cv::Size(), center, maxRadius, flags);
    _warp_polar(src, linearPolar2, center, maxRadius);
    
    std::vector<cv::Point> nonzeros;
    cv::Mat matNonzeros = linearPolar1 != linearPolar2;
    cv::findNonZero(matNonzeros, nonzeros);
    std::cout << "nonzeros size = " << nonzeros.size() << std::endl;
    
    cv::Mat logPolar1, logPolar2;
    cv::warpPolar(src, logPolar1, cv::Size(), center, maxRadius, flags + cv::WARP_POLAR_LOG);
    _warp_polar(src, logPolar2, center, maxRadius, true);
    
    std::vector<cv::Point> nonzeros1;
    cv::Mat matNonzeros1 = logPolar1 != logPolar2;
    cv::findNonZero(matNonzeros1, nonzeros1);
    std::cout << "nonzeros1 size = " << nonzeros1.size() << std::endl;
    
    cv::Mat recoveredLinearPolar1, recoveredLinearPolar2;
    cv::warpPolar(linearPolar1, recoveredLinearPolar1, src.size(), center, maxRadius, flags + cv::WARP_INVERSE_MAP);
    _warp_polar_inverse(linearPolar1, recoveredLinearPolar2, src.cols, src.rows, center, maxRadius);
    
    std::vector<cv::Point> nonzeros2;
    cv::Mat matNonzeros2 = recoveredLinearPolar1 != recoveredLinearPolar2;
    cv::findNonZero(matNonzeros2, nonzeros2);
    std::cout << "nonzeros2 size = " << nonzeros2.size() << std::endl;
    
    cv::Mat recoveredLogPolar1, recoveredLogPolar2;
    cv::warpPolar(logPolar1, recoveredLogPolar1, src.size(), center, maxRadius, flags + cv::WARP_INVERSE_MAP + cv::WARP_POLAR_LOG);
    _warp_polar_inverse(logPolar1, recoveredLogPolar2, src.cols, src.rows, center, maxRadius, true);
    
    std::vector<cv::Point> nonzeros3;
    cv::Mat matNonzeros3 = recoveredLogPolar2 != recoveredLogPolar1;
    cv::findNonZero(matNonzeros3, nonzeros2);
    std::cout << "nonzeros3 size = " << nonzeros3.size() << std::endl;
    
    cv::imshow("src", src);
    cv::imshow("linearPolar1", linearPolar1);
    cv::imshow("linearPolar2", linearPolar2);
    cv::imshow("logPolar1", logPolar1);
    cv::imshow("logPolar2", logPolar2);
    cv::imshow("recoveredLinearPolar1", recoveredLinearPolar1);
    cv::imshow("recoveredLinearPolar2", recoveredLinearPolar2);
    cv::imshow("recoveredLogPolar1", recoveredLogPolar1);
    cv::imshow("recoveredLogPolar2", recoveredLogPolar2);
    cv::waitKey(0);
    
    return 0;
}
